import os
from argparse import ArgumentParser
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
import tensorflow as tf
# import tensorflow_addons as tfa
# import tensorflow_datasets as tfds
# from tensorflow.keras.callbacks import TensorBoard

class KerasDataset:
    
    def __init__(self, dataset_name):
        self.dataset_name = dataset_name
        (
            self.x_train, self.y_train,
            self.x_test, self.y_test
        ) = self.load_data()

    def load_data(self):
        (x_train, y_train), (x_test, y_test) = self.dataset_name.load_data()
        return x_train, y_train, x_test, y_test

    def get_shape(self):
        return self.x_train.shape[1:], self.y_train.shape[1:]

    def preprocess(self):
        (height, width, channels), _ = self.get_shape()
        print(self.get_shape())
        self.x_train = tf.cast(
            self.x_train.reshape((-1, channels, height, width)), dtype=tf.float32) / 255.
        self.x_test = tf.cast(
            self.x_test.reshape((-1, channels, height, width)), dtype=tf.float32) / 255.

    def get_datasets(self):
        self.preprocess()
        return (
            tf.data.Dataset.zip((
                tf.data.Dataset.from_tensor_slices(self.x_train),
                tf.data.Dataset.from_tensor_slices(self.y_train)
            )),
            tf.data.Dataset.zip((
                tf.data.Dataset.from_tensor_slices(self.x_test),
                tf.data.Dataset.from_tensor_slices(self.y_test)
            ))
        )


from originmodel import VisionTransformer

AUTOTUNE = tf.data.experimental.AUTOTUNE

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--logdir", default="logs")
    parser.add_argument("--image-size", default=32, type=int)
    parser.add_argument("--patch-size", default=4, type=int)
    parser.add_argument("--num-layers", default=4, type=int)
    parser.add_argument("--d-model", default=64, type=int)
    parser.add_argument("--num-heads", default=4, type=int)
    parser.add_argument("--mlp-dim", default=128, type=int)
    parser.add_argument("--lr", default=0.00001, type=float)
    parser.add_argument("--weight-decay", default=1e-4, type=float)
    parser.add_argument("--batch-size", default=320, type=int)
    parser.add_argument("--epochs", default=300, type=int)
    args = parser.parse_args()

    # ds = tfds.load("imagenet_resized/32x32", as_supervised=True)
    train_dataset,test_dataset = tf.keras.datasets.cifar100.load_data()
    # train_dataset, test_dataset = KerasDataset(tf.keras.datasets.cifar100).get_datasets()
    # ds_train = (
    #     ds["train"]
    #     .cache()
    #     .shuffle(5 * 320)
    #     .batch(320)
    #     .prefetch(AUTOTUNE)
    # )
    # ds_test = (
    #     ds["validation"]
    #     .cache()
    #     .batch(320)
    #     .prefetch(AUTOTUNE)
    # )
    # ds_train = (
    #     train_dataset
    #     .cache()
    #     .shuffle(5 * 320)
    #     .batch(320)
    #     .prefetch(AUTOTUNE)
    # )
    # ds_test = (
    #     test_dataset
    #     .cache()
    #     .batch(320)
    #     .prefetch(AUTOTUNE)
    # )

    # strategy = tf.distribute.MirroredStrategy()
    # strategy = tf.distribute.OneDeviceStrategy("GPU:0")
    # with strategy.scope():
    VIT_BASE = {
        "image_size":32,
        "patch_size": 16,
        "num_layers": 12,
        "num_classes": 100,
        "d_model": 768,
        "num_heads": 12,
        "mlp_dim": 3072,
        "channels": 3,
        "dropout": 0.1
    }
    model = VisionTransformer(**VIT_BASE)
    model.compile(
        loss=tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True
        ),
        optimizer=tf.train.AdamOptimizer(
            learning_rate=args.lr, beta1=args.weight_decay
        ),
        metrics=["accuracy"],
    )

    model.fit(train_dataset,epochs=1000,batch_size=320)

    # model.fit(
    #     ds_train,
    #     validation_data=ds_test,
    #     epochs=args.epochs,
    #     validation_freq=3
    #     # callbacks=[TensorBoard(log_dir=args.logdir, profile_batch=0),],
    # )
    # model.save_weights(os.path.join(args.logdir, "vit"))
